import math
import torch
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import mindspore as ms
from src.pytorch_init import pytorch_init
from src.mindspore_init import mindspore_init

class compare():
    def __init__(self, ptmodule, msmodule, input_shape, module_type="net", init_mode="ones", input_num=1, print_result="False"):
        self.separator = "-"*50
        self.ptmodule = ptmodule
        self.msmodule = msmodule
        self.input_shape = input_shape
        self.init_mode = init_mode
        self.input_num = input_num
        self.module_type = module_type
        self.print_result = print_result
        self.pt_init = pytorch_init()
        self.ms_init = mindspore_init()
        
    # generate input_data for net
    def gen_input_data(self):
        if self.input_num == 1:
            input_shape = self.input_shape[0]
            input_data = np.random.standard_normal(input_shape).astype(np.float32)
            return input_data
        elif self.input_num == 2:
            input_shape1, input_shape2 = self.input_shape[:]
            input_data1 = np.random.standard_normal(input_shape1).astype(np.float32)
            input_data2 = np.random.standard_normal(input_shape2).astype(np.float32)
            return input_data1, input_data2
            
    def run_msmodule(self, msmodule, input_data):
        x = ms.Tensor(input_data, ms.float32)
        result = msmodule(x)
        return result
    
    def run_msmodule_2ip(self, msmodule, input_data_1, input_data_2):
        x1 = ms.Tensor(input_data_1, ms.float32)
        x2 = ms.Tensor(input_data_2, ms.float32)
        result = msmodule(x1, x2)
        return result
        
    def run_ptmodule(self, ptmodule, input_data):
        x = torch.FloatTensor(input_data)
        result = ptmodule(x)
        return result
    
    def run_ptmodule_2ip(self, ptmodule, input_data_1, input_data_2):
        x1 = torch.FloatTensor(input_data_1)
        x2 = torch.FloatTensor(input_data_2)
        result = ptmodule(x1, x2)
        return result

    # generate init constant for both net
    def gen_init_constant(self):
        if self.init_mode=="ones":
            self.constant = 1
        elif self.init_mode=="random":
            self.constant = np.random.rand()
        print(self.separator)
        print("init condtant is: ", self.constant)
    
    def compare_results(self, result_pt, result_ms):
        result_pt = result_pt.detach().numpy()
        result_ms = result_ms.asnumpy()
        if self.module_type == "loss":
            print(self.separator)
            print("-pt_loss-: ", result_pt)
            print("-ms_loss-: ", result_ms)
        
        elif self.module_type == "net":
            result_pt_shape = result_pt.shape
            result_ms_shape = result_pt.shape
            print(self.separator)
            print("shape of result_pt: ", result_pt_shape)
            print("shape of result_ms: ", result_ms_shape)
            if self.print_result:
                print(self.separator)
                print("-result_pt-: \n", result_pt)
                print("-result_ms-: \n", result_ms)
            if result_pt_shape != result_ms_shape:
                print("output results are in different shapes!")
            else:
                result_pt = result_pt.reshape(1, -1)
                result_ms = result_ms.reshape(1, -1)
                mean_result_pt = np.mean(result_pt)
                mean_result_ms = np.mean(result_ms)
                s = cosine_similarity(result_pt.tolist(), result_ms.tolist())
                print(self.separator)
                print("-similarity-: ", 1 - (abs(mean_result_pt - mean_result_ms)/mean_result_pt))
                print("-cosine similarity-: ", s[0][0])
                print("-variance-: ", np.var(result_ms-result_pt))
                 
    def start_compare(self):
        print(self.separator)
        print("Start Compare...")
        if self.module_type == "net":
            self.gen_init_constant()
            ptmodule = self.pt_init.net_init(self.ptmodule, self.constant)
            msmodule = self.ms_init.net_init(self.msmodule, self.constant)
        elif self.module_type == "loss":
            ptmodule = self.ptmodule
            msmodule = self.msmodule
        if self.input_num == 1:
            input_data = self.gen_input_data()
            result_pt = self.run_ptmodule(ptmodule, input_data)
            result_ms = self.run_msmodule(msmodule, input_data)
        elif self.input_num == 2:
            input_data_1, input_data_2 = self.gen_input_data()            
            result_pt = self.run_ptmodule_2ip(ptmodule, input_data_1, input_data_2)
            result_ms = self.run_msmodule_2ip(msmodule, input_data_1, input_data_2)
        self.compare_results(result_pt, result_ms)

        

        
        